package com.ashes.belife.common.utils;

import cn.hutool.core.util.StrUtil;
import com.ashes.belife.common.exception.AshesException;
import lombok.extern.slf4j.Slf4j;
import net.sourceforge.tess4j.Tesseract;
import net.sourceforge.tess4j.TesseractException;
import net.sourceforge.tess4j.util.LoadLibs;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.rendering.PDFRenderer;

import javax.imageio.ImageIO;
import javax.imageio.ImageReader;
import javax.imageio.stream.ImageInputStream;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.awt.image.RasterFormatException;
import java.io.File;
import java.io.IOException;

@Slf4j
public class ReadTextFromFileRegion {
    /**
     * @Description: ocr按区域识别文件
     **/
    public static String doRecognize(File file, Rectangle rectangle, String dataPath, String formateName, String dpi)
            throws Exception {
        if (!file.exists()) {
            throw new AshesException(0, "待识别文件不存在");
        }
        // OCR 文字识别
        return doOcr(createImg(file, rectangle, formateName), dataPath, dpi);
    }
	  /**
	   * file: 图片文件
	   * rectangle: 图片中需要识别的区域。如new Rectangle(x, y, w, h), 坐标及宽高
	   * dataPath: tessData的路径（文章开头已添加到环境变量，也可传入）
	   * formateName: file文件的格式。如 JPG, PNG, DICOM等, 
	   * dpi: tesseractOcr识别时传入的参数。如：tesseract.setVariable("user_defined_dpi", dpi);
	   */
    public static String doRecognizeOrcOriginResult(File file, Rectangle rectangle, String dataPath, String formateName, String dpi)
            throws Exception {
        if (!file.exists()) {
            throw new AshesException(0, "待识别文件不存在");
        }
        // OCR 文字识别
        return doOcrImpl(createImg(file, rectangle, formateName), dataPath, false, dpi, true);
    }
	/**
	 * 截取的图像缓冲区
	 */
    private static BufferedImage createImg(File file, Rectangle rectangle, String formateName) throws Exception {
        BufferedImage img = readFile(file, rectangle, formateName);
        if(img == null) {
            log.warn("读取图像异常: img == null");
            throw new Exception("文件读取异常, 创建img为空.");
        }
        return img;
    }
    
    /**
     * @Description: 把文件转为BufferedImage对象，并截取指定区域
     **/
    private static BufferedImage readFile(File file, Rectangle region, String formateName) throws IOException {
        ImageInputStream iis = null;
        ImageReader reader = null;
        try {
            // 创建 ImageInputStream 对象
            iis = ImageIO.createImageInputStream(file);
            // 获取 指定 文件的 ImageReader 实例
            reader = ImageIO.getImageReadersByFormatName(formateName).next();
            // 设置解码器
            reader.setInput(iis);
            // 如果需要截取图像，获取完整的BufferedImage，然后截取指定区域
            BufferedImage image = reader.read(0);
            BufferedImage subImage = null;
            try{
                subImage = image.getSubimage(region.x, region.y, region.width, region.height);
            }catch (RasterFormatException e) {
                log.warn("截取图像异常: 识别区域超出边界. err.msg: {}", e.toString());
                return subImage;
            }
            // 释放完整的BufferedImage
            image.flush();
            return subImage;
        } finally {
            // 关闭资源
            if (reader != null) {
                reader.dispose();
            }
            if (iis != null) {
                iis.close();
            }
        }
    }

    /**
     * @Description: ocr按区域识别pdf文件
     **/
    public static String doPdfRecognize(File pdfFile, String dataPath) throws Exception {
        BufferedImage img = readPdfToImage(pdfFile);
        // 按自定义区域截取图像
        BufferedImage subImage = img.getSubimage(460, 170, 320, 100);
        // OCR 文字识别
        return doOcr(subImage, dataPath, "120");
    }

    /**
     * @Description: 把pdf转为图像，renderImageWithDPI方法的第二个参数dpi影响图像分辨率，经测试dpi为300，生成2k分辨率的图像
     **/
    private static BufferedImage readPdfToImage(File pdfFile) throws IOException {
        try (PDDocument document = PDDocument.load(pdfFile)) {
            PDFRenderer renderer = new PDFRenderer(document);
            return renderer.renderImageWithDPI(0, 300);
        }
    }

    /**
     * @Description: 使用Tesseract进行ocr识别
     **/
    private static String doOcr(BufferedImage img, String dataPath, String dpi) throws TesseractException {
        return doOcrImpl(img, dataPath, true, dpi, false);
    }

    /**
     * ocr识别
     * @param img
	   * dataPath: tessData的路径（文章开头已添加到环境变量，也可传入）
     * @param replacedEmp 是否替换回车和空格为空, true:替换, false, 不替换（含回车和空格符）
     * @param dpi 分辨率, 默认 96
     * @param charNoLimit 识别空格: true, 不识别空格: false
     * @return
     * @throws TesseractException
     */
    private static String doOcrImpl(BufferedImage img, String dataPath, boolean replacedEmp, String dpi, boolean charNoLimit) throws TesseractException {
        // 初始化 OCR 引擎
        Tesseract tesseract = new Tesseract();
        File tessDataFolder = LoadLibs.extractTessResources("tessdata");
        tesseract.setDatapath(tessDataFolder.getAbsolutePath());
        //语言包位置  根据实际环境修改替换
        tesseract.setDatapath(dataPath);
        //配置使用的语言  中文
        tesseract.setLanguage("eng");
        if(!charNoLimit) {
            //限制只识别数字字母
            tesseract.setTessVariable("tessedit_char_whitelist", "0123456789CDFGMRTX");
        }
        //设置分辨率
        tesseract.setTessVariable("user_defined_dpi", dpi);
        String result = tesseract.doOCR(img);
        if(replacedEmp) {
            // 文字识别-过滤空白、换行符
            result = result.replace(StrUtil.SPACE, StrUtil.EMPTY).replace(StrUtil.LF, StrUtil.EMPTY);
        }
        return result;
    }
}

